package study.sql

import org.apache.spark.SparkConf
import org.apache.spark.sql.{Encoder, Encoders, SparkSession, functions}
import org.apache.spark.sql.expressions.Aggregator

/**
 * 自定义UDAF聚合函数
 *
 * @author zh
 * @date 2021/5/26 18:08
 */
object CustomUDAF {
  def main(args: Array[String]): Unit = {
    val sparkConf = new SparkConf().setMaster("local").setAppName("sparksql")

    val spark = SparkSession.builder().config(sparkConf).getOrCreate()
    // 隐式转换
    import spark.implicits._

    val dataFrame = spark.read.json("data/user.json")


    // 使用dataSet的方式
    val dataSet = dataFrame.as[User]

    // 构造聚合函数
    val ageAvg = new AgeAvg
    // 将聚合函数转为可查询的列
    val avgColumn = ageAvg.toColumn

    dataSet.select(avgColumn).show()

    // 使用dataFrame的方式
    dataFrame.createOrReplaceTempView("user")

    spark.udf.register("ageAvg",functions.udaf(new AgeAvg1))
    spark.sql("select ageAvg(age) from user").show()




    spark.close()
  }
}

//case class User(age:Long,name:String)

// Aggregator[-IN, BUF, OUT]()
// 三个泛型分别为输入类型，Buf类型，输出类型
class AgeAvg1 extends Aggregator[Long,Buf,Double]{
  // 缓冲区的初始值
  override def zero: Buf = {
    Buf(0,0)

  }

  // 聚合到缓冲区
  override def reduce(b: Buf, a: Long): Buf = {
    b.ageTotal += a
    b.count += 1
    b
  }

  // 缓冲区merge
  override def merge(b1: Buf, b2: Buf): Buf = {
    b1.ageTotal += b2.ageTotal
    b1.count += b2.count
    b1

  }

  // 最终结果
  override def finish(reduction: Buf): Double = {
    reduction.ageTotal.toDouble / reduction.count

  }

  // buffer的编码
  override def bufferEncoder: Encoder[Buf] = {
    Encoders.product
  }

  // 输出的编码
  override def outputEncoder: Encoder[Double] = {
    Encoders.scalaDouble
  }
}

case class Buf(var ageTotal : Long,var count:Long)


class AgeAvg extends Aggregator[User,Buf,Double]{
  // 缓冲区的初始值
  override def zero: Buf = {
    Buf(0,0)

  }

  // 聚合到缓冲区
  override def reduce(b: Buf, a: User): Buf = {
    b.ageTotal += a.age
    b.count += 1
    b
  }

  // 缓冲区merge
  override def merge(b1: Buf, b2: Buf): Buf = {
    b1.ageTotal += b2.ageTotal
    b1.count += b2.count
    b1

  }

  // 最终结果
  override def finish(reduction: Buf): Double = {
    reduction.ageTotal.toDouble / reduction.count

  }

  // buffer的编码
  override def bufferEncoder: Encoder[Buf] = {
    Encoders.product
  }

  // 输出的编码
  override def outputEncoder: Encoder[Double] = {
    Encoders.scalaDouble
  }
}
